// main.js
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamablehttp.js';
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import axios from "axios";
import http from 'node:http';
import { randomUUID } from 'node:crypto';
import 'dotenv/config';

const APIKEY = process.env.API_KEY;
const PORT = Number(process.env.PORT);
console.log(`秘钥: ${APIKEY}`);

// ---------------------------------------------------------------------
// Main server startup
// ---------------------------------------------------------------------
async function main() {
  // 创建 MCP 服务器实例
  const server = new Server(
    { name: 'mcp-http-weather-server', version: '1.0.0' },
    { capabilities: { tools: {} } }
  );

  // 1. 工具列表接口
  server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [
      {
        name: 'getWeather',
        description: '获取指定城市实时天气',
        inputSchema: {
          type: 'object',
          properties: {
            city: { type: 'string', description: '城市名称，如 北京' }
          },
          required: ['city']
        }
      }
    ]
  }));

  // 2. 工具调用接口
  server.setRequestHandler(CallToolRequestSchema, async (request) => {
    const { name, arguments: args } = request.params;
    if (name !== 'getWeather') throw new Error('Unknown tool');

    const { city } = args;
    const apiUrl = "http://apis.juhe.cn/simpleWeather/query";

    try {
      const { status, data } = await axios.get(apiUrl, {
        params: {
          key: APIKEY,
          city: city || "北京",
        }
      });

      if (status !== 200) throw new Error(`API 返回状态: ${status}`);
      
      // 返回格式化后的天气数据
      return {
        content: [{
          type: "text",
          text: JSON.stringify(data, null, 2)
        }],
      };
    } catch (error) {
      console.error('天气查询失败:', error);
      throw new Error(`获取天气数据失败: ${error.message}`);
    }
  });

  // 创建 StreamableHTTP 传输
  const transportOptions = {
    sessionIdGenerator: () => randomUUID(),
    onsessioninitialized: (sessionId) => { // Explicitly type sessionId
      console.log(`MCP Session initialized: ${sessionId}`);
    },
    // enableJsonResponse: true, // You can enable this if you prefer JSON over SSE for non-streaming
  };
  const transport = new StreamableHTTPServerTransport(transportOptions);

  // 启动
  await server.connect(transport);

  console.log(`MCP Server listening on port ${PORT} with transport ${transport.constructor.name}`);

  const httpServer = http.createServer(async (req, res) => {
    // Handle health check endpoint
    if (req.url === '/healthz' && req.method === 'GET') {
      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ status: 'ok', timestamp: new Date().toISOString() }));
      return;
    }

    // MCP requests are POST
    if (req.method === 'POST') {
      let body = '';
      req.on('data', chunk => {
        body += chunk.toString(); // convert Buffer to string
      });
      req.on('end', async () => {
        try {
          const parsedBody = JSON.parse(body);
          // Pass to the transport
          await transport.handleRequest(req, res, parsedBody);
        } catch (e) {
          res.writeHead(400, { 'Content-Type': 'application/json' });
          res.end(JSON.stringify({ error: 'Invalid JSON RPC request: ' + e.message }));
        }
      });
    } else if (req.method === 'GET' && req.headers.accept === 'text/event-stream') {
      // Handle SSE connections if the transport supports them directly for GET
      // (StreamableHTTPServerTransport handles SSE internally for established sessions via POST initially)
      // This explicit GET handler for SSE might be redundant if transport only uses POST for session initiation
      // and then client uses that session ID for a GET SSE stream.
      // For now, assuming transport.handleRequest covers this if a session ID is part of the GET request URL.
      await transport.handleRequest(req, res);
    } else {
      res.writeHead(405, { 'Content-Type': 'text/plain' });
      res.end('Method Not Allowed');
    }
  });

  httpServer.listen(PORT, () => {
    console.log('info', `Chlorpromazine MCP Server listening on HTTP port ${PORT}`);
    console.log('info', `SSE transport expected at /mcp (default path, check SDK docs if different)`);
    console.log('info', `Make sure your MCP client is configured to connect to: http://localhost:${PORT}/mcp`);
    console.log('info', 'Registered MCP Methods:');
    console.log('info', '  - sampling/createMessage (handles models: sober_thinking, fact_checked_answer, buzzkill)');
    console.log('info', '  - tools/call (handles tool: kill_trip)');
    console.log('info', '  - tools/list (lists tool: kill_trip)');
  });

  // Graceful shutdown
  process.on('SIGTERM', () => {
    console.log('info', 'SIGTERM signal received: closing HTTP server');
    httpServer.close(() => {
      console.log('info', 'HTTP server closed');
      server.close(); // Close the MCP server
      console.log('info', 'MCP Server closed');
      process.exit(0);
    });
  });
}

// Initialize Zod schemas for argument validation if not already done
// (They are defined at the top level now, so no action needed here)

main().catch(error => {
  console.log('error', 'Failed to start server', { error: error.message, stack: error.stack });
  process.exit(1);
});